let depthRenderer; // 深度渲染器
let depthMap;
const canvas = document.getElementById("renderCanvas"); // Get the canvas element
// 尝试获取 WebGL2 上下文
const gl = canvas.getContext("webgl2");
if (gl) {
  console.log("支持 WebGL 2.0");
} else {
  alert("设备不支持 WebGL 2.0");
}
// const vertexBuffer = new BABYLON.VertexBuffer()
const engine = new BABYLON.Engine(canvas, true, { alpha: true }); // Generate the BABYLON 3D engine
// engine.setAlphaMode(BABYLON.Engine.ALPHA_PREMULTIPLIED); // 设置透明模式
function showWorldAxis(size, scene) {
  const makeAxis = function (axis, color) {
    const axisMaterial = new BABYLON.StandardMaterial("axisMat", scene);
    axisMaterial.diffuseColor = color;
    axisMaterial.emissiveColor = color;

    const axisLine = BABYLON.MeshBuilder.CreateLines(
      "axis" + axis,
      {
        points: [
          BABYLON.Vector3.Zero(),
          axis === "X"
            ? new BABYLON.Vector3(size, 0, 0)
            : axis === "Y"
            ? new BABYLON.Vector3(0, size, 0)
            : new BABYLON.Vector3(0, 0, size),
        ],
      },
      scene
    );
    axisLine.color = color;
  };

  makeAxis("X", BABYLON.Color3.Red()); // 红色X轴
  makeAxis("Y", BABYLON.Color3.Green()); // 绿色Y轴
  makeAxis("Z", BABYLON.Color3.Blue()); // 蓝色Z轴
}
// 创建后处理
const createPostProcess = (camera) => {
  // 注册 shader 源码
  BABYLON.Effect.ShadersStore["customShaderVertexShader"] = /* glsl */ `
    precision highp float;
    // 必须有这些内置变量
    attribute vec2 position;
    varying vec2 vUV;

    void main(void) {
        vUV = (position + 1.0) * 0.5;
        gl_Position = vec4(position, 0.0, 1.0);
    }
`;

  BABYLON.Effect.ShadersStore["customShaderFragmentShader"] = /* glsl */ `
    precision highp float;
    varying vec2 vUV;
    uniform mat4 invertProjectionMatrix;
    uniform mat4 invertViewMatrix;
    uniform sampler2D textureSampler;
    uniform sampler2D depthMap;
    uniform vec3 sunWorldPos; // 太阳位置
    uniform float time;
    #define PI 3.1415926535
    #define RAYLEIGH_SCALAR_HEIGHT 8500.0 
    #define MIE_SCALAR_HEIGHT 1200.0 
    #define PLANET_RADIUS 6371000.0 
    #define OZONE_CENTER_HEIGHT 30000.0
    #define OZONE_WIDTH 15000.0
    #define G 0.76
    #define ATMOSPHERE_THICKNESS 100000.0 
    #define SAMPLE_1 8
    #define SAMPLE_2 8
    #define SUN_INTENSITY 12.0
    // 计算射线与大气层的交点
   vec2 ray_sphere_intersection(vec3 ray_origin, vec3 ray_direction, vec3 sphere_center, float sphere_radius)
    {
        // 仅考虑大气层内 && 地球表面之上的观察者
        float a = dot(ray_direction, ray_direction);
        vec3 oc = ray_origin - sphere_center;
        float b = 2.0 * dot(ray_direction, oc);
        float c = dot(oc, oc) - (sphere_radius * sphere_radius);
        float d = (b*b) - 4.0*a*c;

        // 返回击中结果，y小于x代表无结果
        if (d < 0.0) return vec2(0.0,0.0);
        // 击中的话有两个相同或者不同的结果
        return vec2(
            (-b - sqrt(d))/(2.0*a),
            (-b + sqrt(d))/(2.0*a)
        );
    }
// 大气密度比例
float DensityRatio(float H,float h){
	return exp(-h/H);
}
 // 瑞丽散射系数
 vec3 RayleighCoefficient(float h){
        const vec3 sigma = vec3(5.802,13.558,33.1) * 1e-6;
        float H_R = RAYLEIGH_SCALAR_HEIGHT;       // 大气层高度
        float rho_h = DensityRatio(H_R,h);
        return sigma * rho_h;
    }
// 米氏散射系数
 vec3 MieCoefficient(float h){
        const vec3 sigma = vec3(33.1,13.5,5.8) * 1e-6;
        float H_M = MIE_SCALAR_HEIGHT;
        float rho_h = DensityRatio(H_M,h);
        return sigma * rho_h;
    }
// 瑞利相位函数
float RayLePhase(float cos_theta){
        return 3.0 / (16.0 * PI) * (1.0 + cos_theta*cos_theta);
    }
// 米氏相位函数 
float MiePhase( float cos_theta){
        float a = 3.0/(8.0*PI);
        float b = (1.0-G*G)/(2.0+G*G);
        float c = 1.0+cos_theta*cos_theta;
        float d = pow(1.0+G*G-2.0*G*cos_theta,1.5);
        return a * b * c / d;
    }
vec3 MieAbsorption(float h){
        const vec3 sigma = vec3(4.4) * 1e-6;
        float H_M = MIE_SCALAR_HEIGHT;
        float rho_h = exp(-(h/H_M));
        return sigma * rho_h;
    }
vec3 OzoneAbsorption(float h){
        const vec3 sigma = vec3(0.650,1.881,0.085) * 1e-6;
        float center = OZONE_CENTER_HEIGHT;
        float width = OZONE_WIDTH;
        float rho_h = max(0.0,1.0-abs((h-center)/width));
        return sigma * rho_h;
    }
vec3 Trasmittance0(vec3 p,vec3 dir,float ds){
		vec3 pp = p;
		float h = length(pp) - PLANET_RADIUS;
		return (RayleighCoefficient(h) + MieCoefficient(h))*ds; 
}
vec3 Trasmittance(vec3 p,vec3 dir,float ds,int SAMPLE_TIME){
		vec3 sum = vec3(0.0);
    vec3 r = vec3(0.0);
    vec3 m = vec3(0.0);
    vec3 absorption = vec3(0.0);
		for(int j=0;j<SAMPLE_TIME;j++){
			vec3 pp = p + dir * ds *( 0.5 + float(j));
			float h = length(pp) - PLANET_RADIUS;
      r += RayleighCoefficient(h);
      m += MieCoefficient(h);
      absorption += MieAbsorption(h) + OzoneAbsorption(h);
			sum = (r + m + absorption)*ds;
		}
		return exp(-sum);
    // return mat3(exp(-r),exp(-m),exp(-sum));
}
vec3 AtmosphereColor(vec3 worldPos,vec3 worldCameraPos,vec3 sunWorldPos){
	
	vec3 viewDir = normalize(worldPos - worldCameraPos);
	vec3 sunDir = normalize(sunWorldPos - worldCameraPos);
	// return vec3(dot(sunDir,viewDir));
	vec3 origin = vec3(0.0);//vec3(worldCameraPos.x,0.0,worldCameraPos.z);//vec3(0.0);
	vec2 n = ray_sphere_intersection(worldCameraPos,viewDir,origin,PLANET_RADIUS+ATMOSPHERE_THICKNESS);
	// 无相交直接返回黑或背景色
    // if (n.y <= 0.0) {
    //     return vec3(1.0,0.0,0.0);
    // }
	float cosTheta = max(dot(viewDir,sunDir),0.0); 
	// 视线与大气层交点
	vec3 A = worldCameraPos;
	vec3 B = max(n.x,n.y) * viewDir + worldCameraPos;
	float ds = max(n.x,n.y)/float(SAMPLE_1);

	vec3 sum = vec3(0.0);	// AB点间的透射率参数总和
	vec3 K = vec3(0.0);		// 积分部分的总计	
	// 进行采样积分
	for(int i=0;i<SAMPLE_1;i++){
		vec3 T1 = vec3(1.0); // 采样点透射率1
		vec3 T2 = vec3(1.0); // 采样点透射率2
		vec3 p = A + ds*viewDir*(float(i)+0.5);
		float h = length(p) - PLANET_RADIUS;
		// pA的透射率
		// sum += Trasmittance0(p,viewDir,ds);
		// T *= exp(-sum);	
    T1 *= Trasmittance(worldCameraPos,viewDir,ds,SAMPLE_2);
		// 光路积分
		vec3 origin = vec3(0.0);
		vec2 n2 = ray_sphere_intersection(p,sunDir,origin,PLANET_RADIUS+ATMOSPHERE_THICKNESS);
		float ds2 = max(n2.x,n2.y)/float(SAMPLE_2);
		T2 *= Trasmittance(p,sunDir,ds2,SAMPLE_2);
		// 计算一些其他项
		vec3 phase = RayLePhase(cosTheta)*RayleighCoefficient(h) +  MiePhase(cosTheta)*MieCoefficient(h);
		K += (T1*T2*phase)*ds;
	}
	return K * SUN_INTENSITY ;
}
    void main(void) {
      vec4 color = texture2D(textureSampler, vUV);
      float depth = texture2D(depthMap, vUV).x;
      float ndcZ = depth * 2.0 - 1.0;
      vec2 ndcXY = vUV * 2.0 - 1.0;  // 将[0,1]转换为[-1,1]
      vec4 viewPos = invertProjectionMatrix * vec4(ndcXY, ndcZ, 1.0);
      viewPos /= viewPos.w;  // 透视除法
      vec4 worldPos = invertViewMatrix * (viewPos);
      vec4 worldCameraPos = (invertViewMatrix * vec4(0.0, 0.0, 0.0, 1.0));
      color.rgb = AtmosphereColor(worldPos.xyz,worldCameraPos.xyz,sunWorldPos);
      float sunColor = max(dot(normalize(worldPos.xyz-worldCameraPos.xyz),normalize(sunWorldPos-worldCameraPos.xyz)),0.0);
      sunColor = pow(sunColor,50.0);
      // color.rgb += sunColor;
      gl_FragColor = vec4(color.rgb,1.0);
    }
`;

  // 创建后处理
  const customPostProcess = new BABYLON.PostProcess(
    "customPostProcess",
    "customShader",
    ["invertProjectionMatrix", "invertViewMatrix", "sunWorldPos", "time"],
    ["textureSampler", "depthMap"], // 屏幕颜色贴图
    1.0,
    camera
  );

  customPostProcess.onApply = (effect) => {
    const depthMap = depthRenderer.getDepthMap();
    const invertViewMatrix = camera.getViewMatrix().clone().invert();
    const invertProjectionMatrix = camera
      .getProjectionMatrix()
      .clone()
      .invert();
    effect.setTexture("depthMap", depthMap);
    effect.setMatrix("invertProjectionMatrix", invertProjectionMatrix);
    effect.setMatrix("invertViewMatrix", invertViewMatrix);
    const time = performance.now() * 1e-4;
    const sunFar = 1e8;
    // const sunPosition = new BABYLON.Vector3(0, -Math.sin(time) * sunFar, Math.cos(time) * sunFar)
    const sunPosition = new BABYLON.Vector3(
      0,
      Math.sin(time) * sunFar,
      Math.cos(time) * sunFar
    );
    // const sunPosition = new BABYLON.Vector3(
    //   0,
    //   Math.sin(Math.PI /4) * sunFar,
    //   Math.cos(Math.PI /4) * sunFar
    // );
    effect.setVector3("sunWorldPos", sunPosition);
    effect.setVector3("time", time);
  };
  return customPostProcess;
};

const createScene = async function () {
  // Creates a basic Babylon Scene object
  const scene = new BABYLON.Scene(engine);
  // const camera = new BABYLON.ArcRotateCamera(
  //   "camera",
  //   0,
  //   0,
  //   1,
  //   new BABYLON.Vector3.Zero(),
  //   scene
  // );
  // // 设置旋转中心
  // camera.setTarget(new BABYLON.Vector3(0, 6371, 0));
  // // 设置世界坐标位置（起点）
  // camera.setPosition(new BABYLON.Vector3(0, 6371.01, 0.1));
  // camera.lowerBetaLimit = 0.01; // 避免完全俯视
  // camera.upperBetaLimit = Math.PI - 0.01; // 避免翻转
  // camera.attachControl(canvas, true);

  // Parameters: name, position, scene
  const camera = new BABYLON.FlyCamera(
    "FlyCamera",
    new BABYLON.Vector3(0, 6371000, 0),
    scene
  );
  // Airplane like rotation, with faster roll correction and banked-turns.
  // Default is 100. A higher number means slower correction.
  camera.maxZ = 1e7;
  camera.rollCorrect = 10;
  // Default is false.
  camera.bankedTurn = true;
  // Defaults to 90° in radians in how far banking will roll the camera.
  camera.bankedTurnLimit = Math.PI / 2;
  // How much of the Yawing (turning) will affect the Rolling (banked-turn.)
  // Less than 1 will reduce the Rolling, and more than 1 will increase it.
  camera.bankedTurnMultiplier = 1;

  // This attaches the camera to the canvas
  camera.attachControl(canvas, true);

  showWorldAxis(5, scene);

  depthRenderer = scene.enableDepthRenderer(camera);
  createPostProcess(camera); // 后处理
  return scene;
};

const scene = await createScene(); //Call the createScene function
// Register a render loop to repeatedly render the scene
engine.runRenderLoop(function () {
  scene.render();
});
// Watch for browser/canvas resize events
window.addEventListener("resize", function () {
  engine.resize();
});
